import click,sys
from collections import defaultdict
from pathlib import Path

genomesize = {}
def dealsize():
    taxSizeList = '/home/zhengzhiqiang/new_pipeline/doc/tax2size.list'
    with open(taxSizeList) as t2s:
        d = {}
        for i in t2s:
            tax, size = i.split()
            d[tax] = int(size)
            genomesize[tax] = size
        return d

def dealname():
    en2cn = '/home/zhengzhiqiang/new_pipeline/doc/en2cn.list'
    with open(en2cn) as ec:
        d = {}
        for i in ec:
            en, cn = i.split('\t')
            x = '_'.join(en.split())
            d[x] = cn.strip()
        return d

@click.command()
@click.option('-ks', '--ksfile', required=True, type=str, help="input kraken speice result file")
@click.option('-kg', '--kgfile', required=True, type=str, help="input kraken genus result file")
@click.option('-info', '--infofile', required=True, type=str, help="input kraken genus result file")
@click.option('-o', '--out', required=True, type=str, help="out nt result file")
@click.option('-nb','--nbb',required=True,help="neibiao.xls")
@click.option('-co','--cofile',required=True,help="coverage_kk.xls")

def main(ksfile, kgfile, infofile, out,nbb,cofile):
    single_Genus = defaultdict(int)
    taxSize = dealsize()
    name = dealname()
    d_coverage=defaultdict(lambda:'0')
    ksCount = defaultdict(lambda: defaultdict(int))
    scaleksCount = defaultdict(lambda: defaultdict(float))
    scaleSum = defaultdict(float)
    ks2Genus = defaultdict(lambda: defaultdict(int))
    Genus = defaultdict(int)
    tax2name = {}
    Info = []
    sampleid = ksfile.strip().split('/')[-2].split('-')[0]


    with open(cofile,encoding='utf-8') as fcoverage:
        for i in fcoverage:
            c_name,c_taxid,c_coverage,c_size=i.strip().split('\t')
            d_coverage[str(c_taxid)] =str(c_coverage)

    with open(infofile, encoding='utf-8') as f:
        for i in f:
            if i.strip().split('\t')[2] == '脑脊液' or i.strip().split('\t')[2] == '血液':
                Info.append(i.split('\t')[0])

    with open(ksfile) as ksf:
        header = ksf.readline().split()
        fCol = 0
        gCol = 0
        sCol = []
        cCol = 0
        g1Col = 0
        s1Col = 0
        for i, e in enumerate(header):
            if e == 'G':
                gCol = i
            elif e.find('S') != -1:
                sCol.append(i)
                if e == 'S':
                    s1Col = i
            elif e == 'Count':
                cCol = i
            elif e == 'F':
                fCol = i
            elif e == 'G1':
                g1Col = i
        for i in ksf:
            info = i.split()
            if info[2] == '-':
                continue
            if sampleid in Info:
                filters = ['Rhizobium', 'phage', 'uncultured', 'Malassezia_restricta', 'Sphingomonas', 'Comamonas', 'Micrococcus', 'Mycobacterium_leprae', \
                            'Agrobacterium', 'Microbacterium', 'Escherichia_virus', 'Burkholderia_vietnamiensis', 'Pandoraea_pnomenusa', 'Acinetobacter_guillouiae', \
                            'Microbacterium_laevaniformans', 'Yarrowia_lipolytica', 'Herbaspirillum','Xanthomonas','Methylobacterium','Cupriavidus','Delftia', \
                            'Enterobius_vermicularis', 'Wuchereria_bancrofti', 'Aspergillus_rambellii']
            else:
                filters = ['Rhizobium', 'phage', 'uncultured', 'Malassezia_restricta', 'Sphingomonas', 'Comamonas', 'Micrococcus', \
                            'Mycobacterium_leprae', 'Agrobacterium', 'Microbacterium', 'Escherichia_virus', 'Burkholderia_vietnamiensis', \
                            'Pandoraea_pnomenusa', 'Acinetobacter_guillouiae', 'Microbacterium_laevaniformans', 'Yarrowia_lipolytica', 'Enterobius_vermicularis', \
                            'Wuchereria_bancrofti', 'Aspergillus_rambellii']
            taxid = ''
            flag = 0
            for j, e in enumerate(sCol):
                if any([i in info[e] for i in filters]):
                    flag = 0
                    break
                if info[e] != '-':
                    flag = 1
                    taxid, *taxname = info[e].split('_')[1:]
                    tax2name[taxid] = '_'.join(taxname)
            if info[g1Col] == 'G1_77643_Mycobacterium_tuberculosis_complex' and info[s1Col] == '-':
                taxid = '1773'
                tax2name[taxid] = 'Mycobacterium_tuberculosis'
                flag = 1
            elif info[g1Col] == 'G1_120793_Mycobacterium_avium_complex_(MAC)' and info[s1Col] == '-':
                taxid = '1764'
                tax2name[taxid] = 'Mycobacterium_avium'
                flag = 1
            if flag == 0:
                continue
            if taxid in taxSize:
                category = info[2]
                ksCount[taxid][f'{category}@{info[fCol]}@{info[gCol]}'] += int(info[cCol])
                scaleksCount[taxid][f'{category}@{info[fCol]}@{info[gCol]}'] = ksCount[taxid][f'{category}@{info[fCol]}@{info[gCol]}'] / taxSize[taxid]
                if tax2name[taxid].find('Erwinia') > -1:
                    pass
                else:
                    scaleSum[category] += int(info[cCol]) / taxSize[taxid]
                ks2Genus[taxid][f'{category}@{info[fCol]}@{info[gCol]}'] = 0
                Genus[category]+=int(info[cCol])
                single_Genus[info[-3]] += int(info[cCol])
            else:
                sys.exit(f'Length of {taxid} is not exists! Please add it`s taxid and length to /home/zhengzhiqiang/new_pipeline/doc/tax2size.list!')

#    Genus = defaultdict(int)
#    Ge = []

    with open(kgfile) as kgf:
        header = kgf.readline().split()
        fCol = 0
        gCol = 0
        cCol = 0
        for i, e in enumerate(header):
            if e == 'G':
                gCol = i
            elif e == 'Count':
                cCol = i
            elif e == 'F':
                fCol = i
        for i in kgf:
            info = i.split()
            if info[2] == '-':
                continue
            category = info[2]
            ctype = f'{category}@{info[fCol]}@{info[gCol]}'
#            if info[gCol] not in Ge:
#                Ge.append(info[gCol])
#                Genus[category] += int(info[cCol])
            for j in ks2Genus:
                if ctype in ks2Genus[j]:
                    ks2Genus[j][ctype] = info[cCol]

    with open(out, 'w', encoding='gbk') as outf:
        outf.write(f'D\tF\tG\tS\tCN_name\tsReads\tgReads\tRelativeAbundance\tGenusRelativeAbundance\tgenomesize\tcoverage(%)\n')
        for k in scaleksCount:
            for k2, v2 in ksCount[k].items():
                info = k2.split('@')
                for i in scaleSum.keys():
                    if i == info[0]:
                        if info[2].find('Erwinia') > -1:
                            scaleRB = '-'
                            GscaleRB = '-'
                        else:
                            #ratio = int(v2) / int(ks2Genus[k][k2])
                            ratio = int(v2) / int(single_Genus[info[2]])
                            scaleRB = f'{scaleksCount[k][k2] / scaleSum[i] * 100 * ratio:.4}'
                            #GscaleRB = f'{int(ks2Genus[k][k2]) / Genus[i] * 100:.4}'
                            GscaleRB = f'{int(single_Genus[info[2]]) / Genus[i] * 100:.4}'
                        flag = 0
                        if tax2name[k] in name :
                            flag = tax2name[k]
                        if flag:
                            if tax2name[k] == 'Erwinia_amylovora':continue
                            #outf.write('\t'.join((i, info[1], info[2], 'S_' + k + '_' + tax2name[k], name[flag], str(v2), str(ks2Genus[k][k2]), str(scaleRB), str(GscaleRB), genomesize[k],d_coverage[str(k)]))+'\n')
                            outf.write('\t'.join((i, info[1], info[2], 'S_' + k + '_' + tax2name[k], name[flag], str(v2), str(single_Genus[info[2]]), str(scaleRB), str(GscaleRB), genomesize[k],d_coverage[str(k)]))+'\n')
                        else:
                            if tax2name[k] == 'Erwinia_amylovora':continue
                            outf.write('\t'.join((i, info[1], info[2], 'S_' + k + '_' + tax2name[k], '-', str(v2), str(single_Genus[info[2]]), str(scaleRB), str(GscaleRB), genomesize[k],d_coverage[str(k)]))+'\n')
                            #outf.write('\t'.join((i, info[1], info[2], 'S_' + k + '_' + tax2name[k], '-', str(v2), str(ks2Genus[k][k2]), str(scaleRB), str(GscaleRB), genomesize[k],d_coverage[str(k)]))+'\n')

    for line in open(nbb,'r'):
        if line.startswith('all_DNA'):
            S1_DNA=line.strip().split('\t')[1]
        if line.startswith('all_CF'):
            S1_CF=line.strip().split('\t')[1]

    with open(out, 'a', encoding='gbk') as outf2:
        outf2.write(f"-\t-\t-\tS_0_内标_DNA\t内标\t{S1_DNA}\t{S1_DNA}\t-\t3833842\t-\n")
        outf2.write(f"-\t-\t-\tS_0_内标_CF\t内标\t{S1_CF}\t{S1_CF}\t-\t2880\t-\n")

if __name__ == "__main__":
    main()
